'''
Version: 2.0
Author: Yue Zhong
Date: 2024-12-31 19:22:08
Description: 
LastEditors: Yue Zhong
LastEditTime: 2025-01-08 21:19:56
'''

import os

DATABASE_PATH   = "./llm_data/LLMConfig/example_dataset.json"
SPACE_PATH      = "./llm_data/LLMConfig/example_space.json"
TASK_PATH       = "./llm_data/LLMConfig/example_task.json"
PROMPT_PATH     = './llm_data/LLMConfig/example_prompt.json'
KNOWLEDGE_PATH  = './llm_data/LLMConfig/example_knowledge.json'

env_path = './llmconfig.env'

dataset_path=   os.path.abspath(os.environ.get("DATABASE_PATH", DATABASE_PATH))
space_path  =   os.path.abspath(os.environ.get("SPACE_PATH", SPACE_PATH))
task_path   =   os.path.abspath(os.environ.get("TASK_PATH", TASK_PATH))
prompt_path =   os.path.abspath(os.environ.get("PROMPT_PATH", PROMPT_PATH))
knowledge_path = os.path.abspath(os.environ.get("KNOWLEDGE_PATH", KNOWLEDGE_PATH))


EMBED_DIM = 1536
TOKEN_LIMIT = 64*1024
TOKEN_COMPLETION_LIMIT = 8*1024
RELAX_TOKEN = 500  # RELAX_TOKEN is the number of tokens to void token limit

_TOKEN_COUNT_FUNC = None

def set_token_count_func(func):
    global _TOKEN_COUNT_FUNC
    _TOKEN_COUNT_FUNC = func


def get_token_count_func():
    global _TOKEN_COUNT_FUNC
    return _TOKEN_COUNT_FUNC

# openai_model_name = "gpt-3.5-turbo" #gpt-3.5-turbo-instruct是InstructGPT的进化，更符合要求，不会很啰嗦
openai_model_name = "deepseek-chat" #deepseek-chat与openai已兼容
zhipu_model_name = "glm-4-flash"

use_open_ai = True

from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_community.chat_models import ChatZhipuAI

if use_open_ai:
    LLM_MODELS = {
        "suggest": lambda: ChatOpenAI(model=openai_model_name, temperature=0),
        "knowledge": lambda: ChatOpenAI(model=openai_model_name),
        "embedding": lambda: OpenAIEmbeddings(),
    }
else:
    # zhipu, glm-4-flashx
    LLM_MODELS = {
        "suggest": lambda: ChatZhipuAI(model_name=zhipu_model_name, temperature=0),
        "knowledge": lambda: ChatZhipuAI(model_name=zhipu_model_name),
        "embedding": lambda: OpenAIEmbeddings(),# ZhipuAIEmbeddings(model="embedding-3" #不可用，需要缴费
                            # dimensions=1024),
    }

DIGIT = 4